#include <unittest/unittest.h>

#include <cusp/gallery/poisson.h>
#include <cusp/gallery/eye.h>
#include <cusp/coo_matrix.h>
#include <cusp/csr_matrix.h>
#include <cusp/dia_matrix.h>
#include <cusp/ell_matrix.h>
#include <cusp/hyb_matrix.h>
#include <cusp/precond/spai.h>
#include <cusp/frobenius_norm.h>

#include <cusp/io/matrix_market.h>

template <typename MemorySpace>
void TestSpaiPatternA(void)
{
    cusp::csr_matrix<int, float, MemorySpace> A;

    cusp::gallery::poisson5pt(A, 10, 10);

    cusp::precond::spai<float, cusp::device_memory> M(A);

    ASSERT_EQUAL(M.matrix.row_offsets, A.row_offsets);
    ASSERT_EQUAL(M.matrix.column_indices, A.column_indices);

    cusp::csr_matrix<int, float, cusp::device_memory> Ref;
    cusp::io::read_matrix_market_file(Ref, "data/spai_pA_reference.mtx");
    ASSERT_EQUAL(M.matrix.row_offsets, Ref.row_offsets);
    ASSERT_EQUAL(M.matrix.column_indices, Ref.column_indices);
    ASSERT_ALMOST_EQUAL(M.matrix.values, Ref.values);

    ASSERT_ALMOST_EQUAL(cusp::frobenius_norm(M.matrix), 2.95968);
}
/* the preconditioner is always on device, anyway */
DECLARE_HOST_DEVICE_UNITTEST(TestSpaiPatternA);

template <typename MemorySpace>
void TestSpaiPatternDiag(void)
{
    cusp::csr_matrix<int, float, MemorySpace> A;
    cusp::csr_matrix<int, float, cusp::host_memory> I;

    cusp::gallery::poisson5pt(A, 10, 10);
    cusp::gallery::eye(I, A.num_rows, A.num_cols);

    cusp::precond::spai<float, cusp::device_memory> M(A, I);

    ASSERT_EQUAL(M.matrix.row_offsets, I.row_offsets);
    ASSERT_EQUAL(M.matrix.column_indices, I.column_indices);

    ASSERT_ALMOST_EQUAL(cusp::frobenius_norm(M.matrix), 2.04348);
}
/* the preconditioner is always on device, anyway */
DECLARE_HOST_DEVICE_UNITTEST(TestSpaiPatternDiag);

template <typename MatrixType>
void TestSpaiFormats(void)
{
    MatrixType A;
    MatrixType I;
    cusp::csr_matrix<int, float, typename MatrixType::memory_space> Icsr;

    cusp::gallery::poisson5pt(A, 10, 10);
    cusp::gallery::eye(I, A.num_rows, A.num_cols);
    Icsr = I;

    cusp::precond::spai<float, cusp::device_memory> M(A, I);

    ASSERT_EQUAL(M.matrix.row_offsets, Icsr.row_offsets);
    ASSERT_EQUAL(M.matrix.column_indices, Icsr.column_indices);

    ASSERT_ALMOST_EQUAL(cusp::frobenius_norm(M.matrix), 2.04348);
}
DECLARE_SPARSE_MATRIX_UNITTEST(TestSpaiFormats);

